import queue
import threading
import subprocess as sp
import websockets
import asyncio
import cv2
import numpy as np
import base64
import time
import random

from algorithm_control.control import Control

import inspect
import ctypes

IMAGE_EXT = [".jpg", ".jpeg", ".webp", ".bmp", ".png"]
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 95]

def _async_raise(tid, exctype):
    """raises the exception, performs cleanup if needed"""
    tid = ctypes.c_long(tid)
    if not inspect.isclass(exctype):
        exctype = type(exctype)
    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
    if res == 0:
        raise ValueError("invalid thread id")
    elif res != 1:
        # """if it returns a number greater than one, you're in trouble,
        # and you should call it again with exc=NULL to revert the effect"""
        ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
        raise SystemError("PyThreadState_SetAsyncExc failed")

def stop_thread(thread):
    _async_raise(thread.ident, SystemExit)


class Live(object):
    def __init__(self):
        self.frame_queue = queue.Queue()
        self.command = ""
        # 自行设置
        self.push_state = False
        self.pullrtsp = ""
        self.pushrtmp = ""
        self.pushflv = ""
        self.aichoice = "face"
        self.aipushstream = False
        self.aipushsubimg = False
        self.control = Control()
        self.loop = asyncio.get_event_loop()

        self.thread_read_frame = threading.Thread(target=Live.read_frame, args=(self,))
        self.thread_push_frame = threading.Thread(target=Live.push_frame, args=(self,))
        self.thread_push_subimg = threading.Thread(target=Live.push_subimg, args=(self,))

    def set_params(self, pullrtsp, aichoice, aipushstream, aipushsubimg):
        if pullrtsp:
            self.pullrtsp = pullrtsp
        self.aichoice = aichoice
        self.aipushstream = aipushstream
        self.aipushsubimg = aipushsubimg
        
        checkcode = ''
        for _ in range(4):
            a = random.randint(0, 9)
            b = str(a)
            checkcode = checkcode + b
        checkcode = 'h264major'
        self.pushrtmp = "rtmp://10.82.14.248:1935/{}/{}".format(self.aichoice, checkcode)
        self.pushflv = "http://10.82.14.248:80/{}/{}.live.flv".format(self.aichoice, checkcode)
        print("设置参数：", self.get_params())

    
    def get_params(self):
        return {"pullrtsp": self.pullrtsp,
                "pushrtmp": self.pushrtmp,
                "pushflv": self.pushflv,
                "aichoice": self.aichoice,
                "aipushstream": self.aipushstream,
                "aipushsubimg": self.aipushsubimg}

    def read_frame(self):
        print("开启拉流线程, 拉流地址: ", self.pullrtsp)
        cap = cv2.VideoCapture(self.pullrtsp)
        # Get video information
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # ffmpeg command
        self.command = ['ffmpeg',
                        '-y',
                        '-f', 'rawvideo',
                        '-vcodec', 'rawvideo',
                        '-pix_fmt', 'bgr24',
                        '-s', "{}x{}".format(width, height),
                        '-r', str(fps),
                        '-i', '-',
                        '-c:v', 'libx264',
                        '-pix_fmt', 'yuv420p',
                        '-preset', 'ultrafast',
                        '-f', 'flv',
                        self.pushrtmp]

        # read webcamera
        while (cap.isOpened()):
            ret, frame = cap.read()

            if not ret:
                print("Opening camera is failed")
                # 说实话这里的break应该替换为：
                # cap = cv2.VideoCapture(self.camera_path)
                # 因为我这俩天遇到的项目里出现断流的毛病
                # 特别是拉取rtmp流的时候！！！！
                break

            # put frame into queue
            
            self.frame_queue.put(frame)

    def push_frame(self):
        # 防止多线程时 command 未被设置
        print("开启推流线程, 推流地址: ", self.pushrtmp, self.pushflv)
        while True:
            if len(self.command) > 0:
                # 管道配置
                print("## 管道配置", self.command)
                self.pipline = sp.Popen(self.command, stdin=sp.PIPE)
                break

        while True:
            if not self.frame_queue.empty():
                frame = self.frame_queue.get()
                # process frame
                # 你处理图片的代码
                # write to pipe
                new_frame = self.control.imageflow(frame)
                # 推ai算法流
                if self.aipushstream:
                    self.pipline.stdin.write(new_frame.tobytes())
                # 推原始流
                else:
                    self.pipline.stdin.write(frame.tobytes())
                # p.stdin.write(frame.tostring())
    
    async def push_subimg_(self):
        print("@@@ 尝试重启推图的websocket, 是否推图:", self.aipushsubimg)
        try:
            async with websockets.connect('ws://10.82.14.248:8000/ws/video/wms/', ping_interval=None) as websocket:
                print("@@@ 推图的websockets 连接成功", self.control.subimg_queue.empty())
                while self.aipushsubimg:
                    if not self.control.subimg_queue.empty():
                        print("推送图片",  self.control.subimg_queue.qsize())
                        subimg = self.control.subimg_queue.get()
                        result, imgencode = cv2.imencode('.png', subimg)
                        data = np.array(imgencode)
                        img = data.tobytes()
                        # base64编码传输
                        img = base64.b64encode(img).decode()
                        print("send img..")
                        await websocket.send(img)
                    time.sleep(1)
        except websockets.exceptions.ConnectionClosedError as e:
            print("@@@ 推图的websocket连接失败")
        except Exception as e:
            print(e)
    
    def push_subimg(self):
        print("开启推图线程")
        while True:
            self.loop.run_until_complete(self.push_subimg_())

    def start_thread(self):
        try:
            print("尝试重启线程")
            self.thread_read_frame.setDaemon(True)
            self.thread_read_frame.start()
            self.thread_push_frame.setDaemon(True)
            self.thread_push_frame.start()
            self.thread_push_subimg.setDaemon(True)
            self.thread_push_subimg.start()
        except:
            print("新建线程")
            self.thread_read_frame = threading.Thread(target=Live.read_frame, args=(self,))
            self.thread_read_frame.setDaemon(True)
            self.thread_read_frame.start()

            self.thread_push_frame = threading.Thread(target=Live.push_frame, args=(self,))
            self.thread_push_frame.setDaemon(True)
            self.thread_push_frame.start()

            self.thread_push_subimg = threading.Thread(target=Live.push_subimg, args=(self))
            self.thread_push_subimg.setDaemon(True)
            self.thread_push_subimg.start()
    
    
    def stop_thread(self):
        while(self.thread_read_frame.is_alive()):
            stop_thread(self.thread_read_frame)
            time.sleep(1)

        while(self.thread_push_frame.is_alive()):
            stop_thread(self.thread_push_frame)
            time.sleep(1)
        
        while(self.thread_push_subimg.is_alive()):
            stop_thread(self.thread_push_subimg)
            time.sleep(1)
        
        try:
            self.pipline.terminate()
            self.pipline.kill()
        except:
            print("kill subprocess fail!")

    def run(self):
        self.stop_thread()
        self.start_thread()